import pandas as pd
import akshare as ak
import time
from datetime import datetime
import tushare as ts

class CompanyDatabase:
    def __init__(self):
        self.cache_time = 3600  # 1小时缓存
        self.last_update = 0
        self.companies = []
    
    def refresh_data(self):
        """从多个数据源获取更完整的数据"""
        if time.time() - self.last_update < self.cache_time:
            return
            
        # 从AKShare获取基础数据
        stock_info = ak.stock_info_a_code_name()
        # 从Tushare获取市值数据
        pro = ts.pro_api()
        df = pro.daily_basic(**{
            'ts_code': ','.join(stock_info['code']),
            'trade_date': datetime.now().strftime('%Y%m%d'),
            'fields': 'ts_code,industry,circ_mv'
        })
        # 合并数据集
        merged_df = pd.merge(stock_info, df, on='ts_code')
        self.companies = merged_df.to_dict('records')
        self.last_update = time.time()
        
    def find_similar(self, industry, market_cap_range):
        """根据行业和市值筛选公司"""
        return [c for c in self.companies 
               if c['industry'] == industry 
               and market_cap_range[0] < c['market_cap'] < market_cap_range[1]] 